import os, argparse, json, logging, subprocess, time, shutil
import yaml
from surfer import Surfer

class TrafficGenerator:
    def __init__(self, config_file:str="config.yml"): 
        self.config_file = config_file

        with open(config_file, "r") as f:
            cfg = yaml.safe_load(f.read())

            # When an item is missing, use default value.
            self.config = {
                "PACKET_CAPTURE_DURATION": cfg.get("PACKET_CAPTURE_DURATION", 10),
                "FILTER_CONDITION"       : cfg.get("FILTER_CONDITION", "doh"),
                "INTERFACE"              : cfg.get("INTERFACE", "any"),
                "DIR_TO_SAVE_TRAFFIC"    : cfg.get("DIR_TO_SAVE_TRAFFIC", "data/traffic/raw"),
                "WAIT_TIME"              : cfg.get("WAIT_TIME", 4),
                "DELAY_TIME"             : cfg.get("DELAY_TIME", 1),
                "PAGE_LOAD_TIMEOUT"      : cfg.get("PAGE_LOAD_TIMEOUT", 5),
                "LOG_FILE"               : cfg.get("LOG_FILE", "EDTG.log")
            }

        if not os.path.exists(self.config["DIR_TO_SAVE_TRAFFIC"]):
            os.makedirs(self.config["DIR_TO_SAVE_TRAFFIC"])

        assert self._check_config() is True

    def _check_config(self) -> bool:
        """Check whether the configuration is valid

        rule_1: PACKET_CAPTURE_DURATION > PAGE_LOAD_TIMEOUT + 2 * DELAY_TIME
        rule_2: PACKET_CAPTURE_DURATION >= PAGE_LOAD_TIMEOUT + DELAY_TIME + WAIT_TIME

        Returns:
            Whether the configuration is valid.
        """

        if self.config["PACKET_CAPTURE_DURATION"] <= \
            2 * self.config["DELAY_TIME"] + self.config["PAGE_LOAD_TIMEOUT"]:
            return False
        
        if self.config["PACKET_CAPTURE_DURATION"] < \
            self.config["PAGE_LOAD_TIMEOUT"] + \
            self.config["DELAY_TIME"] + \
            self.config["WAIT_TIME"]:
            return False

        return True

    def _generate_filepath(self, domain:str, create_subdir:bool=False,
        subdir_root:str="batch", extension:str="pcap", index:int=0,
        status:int=0) -> str:
        """Generate a filename containing its generated time.

        For the case creat_subdir is false, it will return traffic file path under the 
        "DIR_TO_SAVE_TRAFFIC" directly, e.g.
        
        >>> /DIR_TO_SAVE_TRAFFIC/index-domain-time-status.pcap.

        Otherwise, it will return traffic file path as following format:

        >>> /DIR_TO_SAVE_TRAFFIC/subdir_root/domain/index-domain-time-status.pcap

        The status is just a placeholder and it is meaningless when filepath is created,
        as only after the web page is visited, we can get the real status infomation.
        Status indicates whether the web page was successfully accessed, 0 means OK 
        while others means error. 

        Args:
            domain: Domain name, like "apple.com".
            create_subdir: wheter to create subdir.
            subdir_root: Subdir root, only works when `create_subdir` is True.
            extension: Extension of file name.
            index: Index in file name.
            status: Status placeholder, meaningless. 
            
        Returns:
            filename.
        """
        if create_subdir is False:
            return os.path.join(
                self.config["DIR_TO_SAVE_TRAFFIC"], 
                "{index}-{domain}-{time}-{status}.{extension}".format(
                    index = str(index),
                    domain = domain,
                    time = time.strftime("%Y%m%d%H%M%S", time.localtime()),
                    extension = extension,
                    status = status
                )
            )
        
        subdir = os.path.join(self.config["DIR_TO_SAVE_TRAFFIC"], subdir_root, domain)
        if not os.path.exists(subdir):
            os.makedirs(subdir)
        
        return os.path.join(
            self.config["DIR_TO_SAVE_TRAFFIC"],
            subdir_root,
            domain,
            "{index}-{domain}-{time}-{status}.{extension}".format(
                index = str(index),
                domain = domain,
                time = time.strftime("%Y%m%d%H%M%S", time.localtime()),
                extension = extension,
                status = status
            )
        )

    def _replace_status(self, original_file:str, status:int=1) -> str:
        """Replace status of original file and rename original file.

        Args:
            original_file: Original file name, looks like
                'index-domain-time-status.extension',
                e.g. '0-microsoft.com-20220128203615-0.pcap'.
            status: New status.

        Returns: 
            new_file: filename with new status.
        """
        new_file = original_file[:-6] + str(status) + original_file[-5:]
        os.rename(original_file, new_file)

        return new_file

    def _capture(self, duration: int = 10,
        interface: str = "any",
        filter_condition: str = "",
        output_file: str = "out.pcap"):
        """Auto start and stop capture traffic.

        Args:
            duration: Packet capture duration (second).
            output_file: File path to save packets.
        """
        cmd = "tcpdump -G {duration} -W 1 -i {interface} '{filter_condition}' -w {output_file}".\
            format(
                duration = duration,
                interface = interface,
                filter_condition = filter_condition,
                output_file = output_file
            )
        subprocess.Popen(cmd, shell=True) 

    def _get_status_by_filename(self, filename) -> int:
        """get status by filename.
        
        The conventional filename format is: 'index-domain-time-status.extension',
        where status indicates the web page's accessibility, 0 means OK, others means error.

        Args:
            filename: File name.

        Returns:
            Status code, 0 means OK, others means error.
        """
        return int(filename.split("-")[-1].split(".")[0])

    def _remove(self, path_or_dir:str):
        """Remove file or directory.
        """
        if os.path.isfile(path_or_dir):
            os.remove(path_or_dir)
        else:
            shutil.rmtree(path_or_dir)

    def run(self, domain:str, create_subdir:bool=False, subdir_root:str="", index:int=0) -> str:
        """Generate DoH traffic of visiting single domain.

        Workflow:
        1. Start capturing packet
        2. Sleep DELAY_TIME
        3. Visit website (timeout: PAGE_LOAD_TIMEOUT)
        4. Sleep DELAY_TIME
        5. Change status bit in pcap file names.

        The whole process will take around 

        Args:
            domain: Domain name, or URL.
            create_subdir: wheter to create subdir.
            subdir_root: Subdir root, only works while `create_subdir` is True.
            index: Index in output file name.

        Returns:
            output_file: Output pcap file path.
        """

        # Generate output file dir and path
        output_file = self._generate_filepath(
            domain = domain,
            create_subdir = create_subdir,
            subdir_root = subdir_root,
            index = index
        )

        # Packet capture
        self._capture(
            duration = self.config["PACKET_CAPTURE_DURATION"],
            interface = self.config["INTERFACE"],
            filter_condition = self.config["FILTER_CONDITION"],
            output_file = output_file
        )

        # Delay
        time.sleep(self.config["DELAY_TIME"])

        # Visit website
        surfer = Surfer(self.config_file)
        status, info = surfer.visit(domain)
        time.sleep(self.config["WAIT_TIME"])
        surfer.quit()

        # Replace the status in the filename
        if status != 0:
            output_file = self._replace_status(output_file)
            print(info)

        return output_file

    def run_batch(self, top:int, amount:int, domains_file:str) -> list:
        """Generate DoH traffic by visiting a list of domains form domain list.

        For a domain name, it will generate several DoH traffic trace
        (The specific number is defined by arguement 'amount'). If generator
        failed max_failed_times  when visiting domain d_i, d_i will be regarded 
        as "bad domain", and related traffic will be removed. This operation
        guarantee target domain's access success rate is greater than 50%.

        max_failed_times = amount if amount <= 5 else max(5, int(amount * 0.3))

        Args:
            top: Top n domains to read.
            amount: The amount a domain name is visited repeatedly.
            domains_file: Domain file path.

        Returns:
            failed_domains: List of filed domains.
        """
        logging.basicConfig(
                    filename=self.config["LOG_FILE"],
                    format="%(asctime)s - %(message)s",
                    level=logging.INFO
                )
        logging.info("Start. domain_file:{}  top:{}  amount:{}".format(domains_file, top, amount))
        failed_domains, succeed_domains, domains = [], [], []
        max_failed_times = amount if amount <= 5 else max(5, int(amount * 0.3))
        
        # Gets the number of lines of the domains file.
        lines = 0

        with open(domains_file) as f_in:
            for _ in f_in.readlines():
                lines += 1

        assert top <= lines

        # if top == 0, load all domain names.
        if top <= 0:
            top = lines

        with open(domains_file) as f_in:
            for domain in f_in.readlines():               
                domains.append(domain.strip())

        subdir_root = time.strftime("%Y%m%d%H%M%S", time.localtime())

        for i in range(top):
            j, failed_times = 0, 0

            while j < amount:

                # generate traffic
                filepath = self.run(
                    domain = domains[i],
                    create_subdir = True,
                    subdir_root = subdir_root,
                    index = j
                )
                
                # check traffic
                if self._get_status_by_filename(filepath) != 0:
                    # add failed_times
                    failed_times += 1
                
                    if failed_times == max_failed_times:
                        logging.info("Failed: {}".format(domains[i]))
                        # remove all files of this domain
                        self._remove(os.path.dirname(filepath))
                        failed_domains.append(domains[i])
                        break
                    
                    self._remove(filepath)
                    continue

                time.sleep(self.config["DELAY_TIME"])
                j += 1

            if failed_times < max_failed_times:
                logging.info("Succeed: {}".format(domains[i]))
                succeed_domains.append(domains[i])

        logging.info("End.")
        return failed_domains

def main():
    parser = argparse.ArgumentParser(description="Generate DoH traffic.")

    parser.add_argument(
        "--config", "-c",
        type = str,
        default = "config.yml",
        help = "Config file path. (Default 'config.yml')")

    parser.add_argument(
        "--domain", "-d",
        type = str,
        help = "Single domain. Generate traffic by visiting this domain.")

    parser.add_argument(
        "--domain_list", "-dl",
        type = str,
        help = "Domain list file. Generate traffic by visiting domain(s) from this file."
    )

    parser.add_argument(
        "-t", "--top",
        type = int,
        default = 0,
        help = "Top n domains to read. Default 0, which means all. Only works in '-dl' mode."
    )

    parser.add_argument(
        "-a", "--amount", 
        type = int,
        default = 1,
        help = "The amount of a domain to visit when using domain list file. \
            Default 1. Only works in '-dl' mode."
    )

    args = parser.parse_args()

    with open(args.config, "r") as f:
        cfg = yaml.safe_load(f.read())

    generator = TrafficGenerator(args.config)

    # Single
    if args.domain is not None:
        print("Setting:")
        print(json.dumps(cfg, indent=4))
    
        generator.run(args.domain)

    # Batch 
    elif args.domain_list is not None:
        generator.run_batch(args.top, args.amount, args.domain_list)

    # Error
    else:
        print("Please specify -d option or -dl option.")

def test_remove():
    g = TrafficGenerator()
    os.makedirs("test_dir")
    with open("test_dir/example1.txt","w+") as f:
        f.write("test")
    with open("test_dir/example2.txt","w+") as f:
        f.write("test")
    g._remove("test_dir/example1.txt")
    g._remove(os.path.dirname("test_dir/example2.txt"))

def test_get_status_by_filename():
    g = TrafficGenerator()
    original_file = "0-microsoft.com-20220128203615-0.pcap"
    assert g._get_status_by_filename(original_file) == 0

def test_replace_status():
    g = TrafficGenerator()
    original_file = "0-microsoft.com-20220128203615-0.pcap"

    subprocess.run(["touch", original_file])
    new_file = g._replace_status(original_file)
    print(new_file)
    time.sleep(2)

    subprocess.run(["rm", new_file])

if __name__ == "__main__":
    main()
    # test_replace_status()
    # test_get_status_by_filename()
    # test_remove()